

# 导入OpenAI库
from openai import OpenAI
# 从fastmcp库导入Client类，用于创建MCP客户端
from fastmcp import Client
# 导入asyncio库，用于异步操作
import asyncio
# 导入json库，用于处理JSON数据
import json
# 导入系统路径处理
import sys
# 导入os库，用于文件操作
import os

# 将项目根目录添加到Python路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

# 导入日志记录器
from core.logger import get_logger

# 创建日志记录器实例
logger = get_logger(module_name="LLM_API", log_level="info")

# 调用LLM API
async def call_LLM_api(api_key, messages, model, stream=True, tools=None, base_url=None):
    logger.info(f"开始调用LLM接口 - model:{model}, stream:{stream}")

    # 创建OpenAI客户端
    client = OpenAI(
        api_key=api_key,
        base_url=base_url,  # 使用传入的 base_url
    )
    
    # 构建基本参数字典
    params = {
        "model": model,
        "messages": messages,
        "stream": stream,
        "tools": tools,
    }
    
    # 根据stream参数设置不同的extra_body
    if stream:
        params["stream_options"] = {"include_usage": True}  # 流式输出时，要求模型返回包含usage信息，否则无法计算Token成本
        params["extra_body"] = {"enable_thinking": True}    # Qwen3模型通过enable_thinking参数控制思考过程（开源版默认True，商业版默认False）
    else:
        params["extra_body"] = {"enable_thinking": False}   # Qwen3模型通过enable_thinking参数控制思考过程（开源版默认True，商业版默认False）
    
    logger.debug(f"API参数配置完成: stream={stream}, tools_count={len(tools) if tools else 0}")
    
    try:
        # 使用解包参数调用API
        completion = client.chat.completions.create(**params)
        
        # 根据流式输出与否，选择不同的处理函数
        if stream:
            # 流式输出，等待流式输出完成
            return await process_completion_stream(completion)
        else:
            # 非流式直接返回
            return process_completion(completion)
    except Exception as e:
        logger.error(f"调用LLM API失败: {str(e)}", exc_info=True)
        raise


async def process_completion_stream(completion):
    """
    处理流式输出的完成对象，提取回复内容、推理过程、工具调用等信息
    """
    logger.info("开始处理流式输出")
    reasoning_content = ""    # 深度思考内容
    full_content = ""         # 回复内容
    tool_calls_by_index = {}  # 用字典按index保存工具调用
    completion_tokens = 0     # 回复 Token 长度
    prompt_tokens = 0         # 输入 Token 长度
    finish_reason = None      # 完成原因

    # 用于跟踪是否打印过标题
    printed_reasoning_header = False
    printed_content_header = False

    # 异步处理每个chunk
    for chunk in completion:
        # 检查choices列表是否有元素
        if chunk.choices and len(chunk.choices) > 0:
            delta = chunk.choices[0].delta
            
            # 获取深度思考内容
            if hasattr(delta, 'reasoning_content') and delta.reasoning_content is not None:
                reasoning_content += delta.reasoning_content    # 片段深度思考内容追加到全局变量
                
                # 第一次打印时显示标题
                if not printed_reasoning_header:
                    print("\033[1;36m┌── 深度思考过程 ───────────────────┐\033[0m")
                    printed_reasoning_header = True
                    logger.debug("开始输出模型思考过程")
                
                # 实时打印，不换行
                print(delta.reasoning_content, end="", flush=True)
            
            # 获取回复内容
            if hasattr(delta, 'content') and delta.content is not None:
                full_content += delta.content    # 片段回复内容追加到全局变量
                
                # 第一次打印时显示标题
                if not printed_content_header:
                    if printed_reasoning_header:
                        print("\n\033[1;36m└────────────────────────────────┘\033[0m")
                    print("\033[1;32m┌── 回复内容 ─────────────────────┐\033[0m")
                    printed_content_header = True
                    logger.debug("开始输出模型回复内容")
                
                # 实时打印，不换行
                print(delta.content, end="", flush=True)
            
            # 获取工具调用（根据index处理）
            if hasattr(delta, 'tool_calls') and delta.tool_calls is not None:
                for tool_call in delta.tool_calls:
                    # 使用index作为主键
                    index = tool_call.index
                    
                    # 如果是新的工具调用index，初始化记录
                    if index not in tool_calls_by_index:
                        tool_calls_by_index[index] = {
                            "id": tool_call.id,
                            "type": tool_call.type,
                            "function": {
                                "name": tool_call.function.name if tool_call.function else None,
                                "arguments": ""
                            },
                            "index": index
                        }
                        logger.debug(f"新增工具调用 #{index}: {tool_call.function.name if tool_call.function else '未知'}")
                    else:
                        # 如果已有此index但id为空，现在有了，则更新id
                        if tool_call.id and not tool_calls_by_index[index]["id"]:
                            tool_calls_by_index[index]["id"] = tool_call.id
                            
                        # 如果已有此index但type为空，现在有了，则更新type
                        if tool_call.type and not tool_calls_by_index[index]["type"]:
                            tool_calls_by_index[index]["type"] = tool_call.type
                            
                        # 如果已有此index但name为空，现在有了，则更新name
                        if tool_call.function and tool_call.function.name and not tool_calls_by_index[index]["function"]["name"]:
                            tool_calls_by_index[index]["function"]["name"] = tool_call.function.name
                    
                    # 始终累加arguments (如果有)
                    if tool_call.function and tool_call.function.arguments:
                        tool_calls_by_index[index]["function"]["arguments"] += tool_call.function.arguments
            
            # 获取完成原因
            if hasattr(chunk.choices[0], 'finish_reason') and chunk.choices[0].finish_reason:
                finish_reason = chunk.choices[0].finish_reason
                logger.debug(f"响应完成原因: {finish_reason}")
        
        # 获取词元统计
        if chunk.usage is not None:
            completion_tokens = chunk.usage.completion_tokens
            prompt_tokens = chunk.usage.prompt_tokens

    # 流式输出已完成，转换工具调用字典为列表
    tool_calls = list(tool_calls_by_index.values())
    
    if printed_content_header:
        print("\n\033[1;32m└────────────────────────────────┘\033[0m")
    elif printed_reasoning_header:  # 如果只有思考内容，没有回复内容，也需要关闭思考框
        print("\n\033[1;36m└────────────────────────────────┘\033[0m")
    
    if tool_calls:
        print("\033[1;33m┌── 工具调用 ─────────────────────┐\033[0m")
        for i, tool in enumerate(tool_calls):
            print(f"\033[1;33m│ #{i+1}\033[0m {tool['function']['name']}")
        print("\033[1;33m└────────────────────────────────┘\033[0m")
        logger.info(f"检测到 {len(tool_calls)} 个工具调用")
    
    logger.info(f"Token统计 - 输入: {prompt_tokens}, 输出: {completion_tokens}")
    
    
    # 直接返回结果
    return {
        "reasoning_content": reasoning_content,
        "full_content": full_content,
        "tool_calls": tool_calls,
        "completion_tokens": completion_tokens,
        "prompt_tokens": prompt_tokens,
        "finish_reason": finish_reason
    }


def process_completion(completion):
    """
    处理非流式输出的完成对象，提取回复内容、推理过程、工具调用等信息
    """
    logger.info("处理非流式输出响应")
    # 从API响应对象中提取消息内容
    message = completion.choices[0].message
    
    # 提取推理过程内容
    reasoning_content = message.reasoning_content
    # 提取AI生成的回复内容
    content = message.content
    # 提取AI调用的工具列表
    tool_calls = message.tool_calls
    
    # 获取Token使用统计信息
    usage = completion.usage
    completion_tokens = usage.completion_tokens  # 输出token数量
    prompt_tokens = usage.prompt_tokens          # 输入token数量

    # 获取完成原因
    finish_reason = completion.choices[0].finish_reason
    logger.debug(f"响应完成原因: {finish_reason}")
    
    # 将提取的内容整合为字典返回
    return {
        "reasoning_content": reasoning_content,  # 推理过程
        "full_content": content,                 # 回复内容
        "tool_calls": tool_calls,                # 工具调用列表
        "completion_tokens": completion_tokens,  # 输出token计数
        "prompt_tokens": prompt_tokens,          # 输入token计数
        "finish_reason": finish_reason,          # 完成原因
        "message": message                       # 原始消息对象
    }


# 获取MCP工具并转换为OpenAI格式
async def get_and_convert_tools(client) -> list:
    """
    获取MCP工具并转换为OpenAI格式
    """
    logger.info("开始获取MCP工具列表")
    try:
        # 获取MCP工具列表
        raw_tools = await client.list_tools()
        logger.info(f"成功获取 {len(raw_tools)} 个MCP工具")
        
        # 将MCP工具列表转换为OpenAI格式
        formatted_tools = [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "input_schema": tool.inputSchema
                }
            }
            for tool in raw_tools
        ]
        
        logger.debug(f"工具转换完成，工具列表: {[tool['function']['name'] for tool in formatted_tools]}")
        return formatted_tools
    except Exception as e:
        error_msg = f"获取MCP工具时出错: {str(e)}"
        logger.error(error_msg, exc_info=True)
        raise


# 工具调用LLM API
async def call_LLM_Tool_api(api_key, messages, model, stream=True, mcp_config=None, base_url=None):
    logger.info(f"开始工具增强对话 - model:{model}")
    # 创建MCP客户端实例
    mcp_client = Client(mcp_config)
    
    # 初始化统计变量
    total_prompt_tokens = 0         # 总输入Token
    total_completion_tokens = 0     # 总输出Token
    api_call_count = 0              # API调用次数
    
    try:
        async with mcp_client:
            # 获取并转换工具格式
            formatted_tools = await get_and_convert_tools(mcp_client)
            
            # 持续处理工具调用，直到没有工具调用或处理完成
            iteration = 0
            max_iterations = 100  # 设置最大尝试次数
            
            # 首次调用LLM API
            logger.info("开始第 1 次API调用")
            api_call_count += 1
            response = await call_LLM_api(api_key, messages, model, stream, tools=formatted_tools, base_url=base_url)
            
            # 累计Token使用量
            total_prompt_tokens += response.get('prompt_tokens', 0)
            total_completion_tokens += response.get('completion_tokens', 0)
            
            while True:
                iteration += 1
                
                # 检查响应是否需要调用工具
                if not response.get('tool_calls') or response.get('finish_reason') != 'tool_calls':
                    logger.info("对话已完成，不需要工具调用")
                    
                    # 返回结果
                    logger.info(f"会话统计 - API调用次数: {api_call_count}, 总输入Token: {total_prompt_tokens}, 总输出Token: {total_completion_tokens}")
                    
                    return {
                        "reasoning_content": response['reasoning_content'],     # 深度思考内容
                        "full_content": response['full_content'],               # 回复内容
                        "tool_calls": response['tool_calls'],                   # 工具调用列表
                        "completion_tokens": response['completion_tokens'],     # 输出token计数
                        "prompt_tokens": response['prompt_tokens'],             # 输入token计数
                        "total_prompt_tokens": total_prompt_tokens,             # 总输入token计数
                        "total_completion_tokens": total_completion_tokens,     # 总输出token计数
                        "api_call_count": api_call_count,                       # API调用次数
                    }
                
                # 检查是否超出最大尝试次数
                if iteration > max_iterations:
                    logger.warning(f"已达到最大工具调用尝试次数 ({max_iterations})，停止调用")
                    
                    # 添加系统消息，告诉模型不再尝试调用工具
                    messages.append({
                        "role": "system",
                        "content": f"工具调用次数已达上限({max_iterations}次)，请直接回答用户问题，不要再尝试调用工具。"
                    })
                    
                    # 最后一次调用API，传入终止工具调用的指示，不传入工具列表
                    logger.info("开始最终API调用(无工具模式)")
                    api_call_count += 1
                    response = await call_LLM_api(api_key, messages, model, stream, tools=None, base_url=base_url)
                    
                    # 累计Token使用量
                    total_prompt_tokens += response.get('prompt_tokens', 0)
                    total_completion_tokens += response.get('completion_tokens', 0)
                    
                    logger.info(f"会话统计 - API调用次数: {api_call_count}, 总输入Token: {total_prompt_tokens}, 总输出Token: {total_completion_tokens}")
                    
                    return {
                        "full_content": response['full_content'],               # 回复内容
                        "reasoning_content": response['reasoning_content'],     # 深度思考内容
                        "tool_calls": response['tool_calls'],                   # 工具调用列表
                        "completion_tokens": response['completion_tokens'],     # 输出token计数
                        "prompt_tokens": response['prompt_tokens'],             # 输入token计数
                        "total_prompt_tokens": total_prompt_tokens,             # 总输入token计数
                        "total_completion_tokens": total_completion_tokens,     # 总输出token计数
                        "api_call_count": api_call_count,                       # API调用次数
                    }
                
                logger.info(f"第 {iteration}/{max_iterations} 轮工具调用，需调用 {len(response['tool_calls'])} 个工具")

                # 构造原始模型响应内容，添加到对话历史
                messages.append({
                    "role": "assistant",                                    
                    "content": response['full_content'],                   
                    "tool_calls": [                                         
                        {
                            "id": tool.id if hasattr(tool, 'id') else tool['id'],
                            "function": {                                   
                                "name": tool.function.name if hasattr(tool, 'function') else tool['function']['name'],
                                "arguments": tool.function.arguments if hasattr(tool, 'function') else tool['function']['arguments']
                            },
                            "type": "function"                              
                        } for tool in response['tool_calls']                
                    ]
                })

                # 遍历所有工具调用请求
                for i, tool in enumerate(response['tool_calls']):
                    try:
                        # 根据对象类型使用不同的访问方式获取工具名称
                        tool_name = tool.function.name if hasattr(tool, 'function') else tool['function']['name']
                        logger.info(f"调用工具: {tool_name}")
                        
                        # 调用MCP客户端执行工具，传入工具名称和解析后的参数
                        tool_arguments = json.loads(tool.function.arguments if hasattr(tool, 'function') else tool['function']['arguments'])
                        tool_result = await mcp_client.call_tool(tool_name, tool_arguments)
                        
                        # 提取工具返回内容
                        tool_content = tool_result.content[0].text if tool_result.content else ''
                        if tool_content:
                            preview = tool_content[:50] + ('...' if len(tool_content) > 50 else '')
                            logger.debug(f"工具 {tool_name} 返回结果: {preview}")
                        
                        # 添加工具调用结果
                        tool_id = tool.id if hasattr(tool, 'id') else tool['id']
                        messages.append({
                            "role": "tool",
                            "content": tool_content,
                            "tool_call_id": tool_id
                        })
                    except Exception as e:
                        # 捕获工具调用错误
                        error_message = str(e)
                        # 同样需要安全地获取工具名称
                        tool_name = tool.function.name if hasattr(tool, 'function') else tool['function']['name']
                        logger.error(f"工具 {tool_name} 调用失败: {error_message}", exc_info=True)
                        
                        # 添加工具调用失败结果
                        tool_id = tool.id if hasattr(tool, 'id') else tool['id']
                        messages.append({
                            "role": "tool",
                            "content": f"工具调用出错: {error_message}",
                            "tool_call_id": tool_id
                        })

                # 工具调用完成后再次调用API
                logger.info(f"开始第 {api_call_count + 1} 次API调用")
                api_call_count += 1
                response = await call_LLM_api(api_key, messages, model, stream, tools=formatted_tools, base_url=base_url)
                
                # 累计Token使用量
                total_prompt_tokens += response.get('prompt_tokens', 0)
                total_completion_tokens += response.get('completion_tokens', 0)
    except Exception as e:
        logger.error(f"工具对话会话执行出错: {str(e)}", exc_info=True)
        raise


# 程序入口点
if __name__ == '__main__':
    logger.info("LLM API工具启动")
    print("\033[1;36m══════════════════════════════════\033[0m")
    print("\033[1;36m       工具增强型LLM对话系统       \033[0m")
    print("\033[1;36m══════════════════════════════════\033[0m")
    
    # MCP多服务器配置字典
    mcp_config = {
    "mcpServers": 
        {
            "modelscope_cantian-ai_Bazi-MCP": 
            {
                "url": "https://mcp.api-inference.modelscope.net/8d92e83e08df46/sse",   # Seatable服务端点
            },

            "GYCDWcq578JMgRTd34B0_": 
            {
            "name": "自定义seatable MCP 服务器",
            "type": "stdio",
            "description": "",
            "isActive": True,
            "registryUrl": "",
            "command": "uv",
            "args": [
                "--directory",
                "F:\流量工作室\python_code\WeChat_AI_Assistant_1.1",
                "run", "tools/SeaTable_server.py"
            ]
            }
        }
    }

    #千问
    api_key="sk-4125bc9698694675a5300c96202218e0"
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    model="qwen3-235b-a22b"

    # # deepseek
    # api_key="sk-2dc3054897024436ad47912a07ef4617"
    # base_url="https://api.deepseek.com"
    # model="deepseek-reasoner"

    # 消息列表
    messages=[
        {"role": "system", "content": "你是一名命理师，善于调用工具帮我解决问题。"},
        {"role": "user", "content": "帮我算算我的八字，杨洋，男，农历1994年5月20日，17点30分，出生地：贵州，我想知道我今年跟朋友合伙开公司怎么样"}
    ]

    stream=True

    # 打印开始调用API信息
    print("\033[1;36m[API]\033[0m 开始调用LLM接口...")
    print("\033[1;36m[工具]\033[0m 开始获取MCP工具列表...")

    try:
        # 工具调用LLM API
        print("\033[1;36m┌── 工具对话会话 ─────────────────┐\033[0m")
        
        result = asyncio.run(call_LLM_Tool_api(api_key, messages, model, stream, mcp_config, base_url))
        
        # 如果是非流式输出，需要显示最终结果，因为没有实时输出
        if not stream:
            # 显示思考过程（如果有）
            if result['reasoning_content']:
                print("\033[1;36m┌── 深度思考过程 ───────────────────┐\033[0m")
                print(result['reasoning_content'])
                print("\033[1;36m└────────────────────────────────┘\033[0m")
            
            # 显示回复内容
            print("\033[1;32m┌── 回复内容 ─────────────────────┐\033[0m")
            print(result['full_content'])
            print("\033[1;32m└────────────────────────────────┘\033[0m")
            
            # 显示工具调用（如果有）
            if result['tool_calls']:
                print("\033[1;33m┌── 工具调用 ─────────────────────┐\033[0m")
                for i, tool in enumerate(result['tool_calls']):
                    tool_name = tool.function.name if hasattr(tool, 'function') else tool['function']['name']
                    print(f"\033[1;33m│ #{i+1}\033[0m {tool_name}")
                print("\033[1;33m└────────────────────────────────┘\033[0m")
        
        print("\033[1;36m└────────────────────────────────┘\033[0m")
        
        # 打印会话统计
        print("\033[1;33m┌── 会话统计 ─────────────────────┐\033[0m")
        print(f"\033[1;33m│\033[0m API调用次数: {result['api_call_count']}")
        print(f"\033[1;33m│\033[0m 总输入Token: {result['total_prompt_tokens']}")
        print(f"\033[1;33m│\033[0m 总输出Token: {result['total_completion_tokens']}")
        print("\033[1;33m└────────────────────────────────┘\033[0m")

        # 计算实际费用 (价格为每千Token)
        prompt_cost = (result["total_prompt_tokens"] / 1000) * 0.02
        completion_cost = (result["total_completion_tokens"] / 1000) * 0.002
        total_cost = prompt_cost + completion_cost

        print("\033[1;33m┌── 成本统计 ─────────────────────┐\033[0m")
        print(f"\033[1;33m│\033[0m 输入Token成本: ¥{prompt_cost:.6f}")
        print(f"\033[1;33m│\033[0m 输出Token成本: ¥{completion_cost:.6f}")
        print(f"\033[1;33m│\033[0m 总成本: ¥{total_cost:.6f}")
        print("\033[1;33m└────────────────────────────────┘\033[0m")

        logger.info(f"成本统计 - 输入: ¥{prompt_cost:.6f}, 输出: ¥{completion_cost:.6f}, 总计: ¥{total_cost:.6f}")
        logger.info("工具增强型LLM对话系统正常结束")
    except Exception as e:
        logger.error(f"程序执行出现错误: {str(e)}", exc_info=True)
        print(f"\033[1;31m[错误]\033[0m 程序执行出现错误: {str(e)}")








    








